from moviepy.video.io.VideoFileClip import VideoFileClip
from moviepy.video.compositing.CompositeVideoClip import concatenate_videoclips
import os
import sys

def extract_clips_with_audio(input_video, output_dir, segments):
    """
    从原视频中提取多个片段，保留原有音频，并可选择拼接成完整视频
    
    参数:
        input_video: 输入视频路径
        output_dir: 片段保存目录
        segments: 片段区间列表，格式为 [(start_sec, end_sec), ...] (秒为单位)
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 加载原视频（自动包含音频轨道）
    video = VideoFileClip(input_video)
    clips = []  # 存储提取的片段
    clip_paths = [] # 存储片段路径
    final_path = os.path.join(output_dir, "final_video.mp4")
    
    # 提取每个片段
    for i, (start, end) in enumerate(segments):
        # 裁剪片段（自动保留对应区间的音频）
        clip = video.subclipped(start, end)
        
        # 保存单个片段（含音频）
        clip_path = os.path.join(output_dir, f"clip_{i+1}.mp4")
        clip.write_videofile(
            clip_path,
            codec="libx264",  # 视频编码器（确保兼容性）
            audio_codec="aac",  # 音频编码器
            fps=video.fps  # 保持原帧率
        )
        print(f"已保存片段 {i+1}: {clip_path} (时长: {end-start:.1f}秒)")
        clips.append(clip)
        clip_paths.append(clip_path)
    
    # 拼接所有片段（保持原音频顺序）
    if len(clips) > 1:
        final_clip = concatenate_videoclips(clips, method="compose")
        final_clip.write_videofile(
            final_path,
            codec="libx264",
            audio_codec="aac",
            fps=video.fps
        )
        print(f"已拼接完整视频: {final_path}")
    elif len(clips) == 1:
        os.rename(clip_paths[0], final_path)
        print(f"只有一个片段，已保存完整视频: {final_path}")
    else:
        print("没有需要处理的片段")
    # 释放资源
    video.close()
    for clip in clips:
        clip.close()

if __name__ == "__main__":
    if len(sys.argv) < 5:
        print("用法: python clip_with_audio.py <输入视频> <输出目录> <起始秒> <结束秒>")
        sys.exit(1)
    
    input_video = sys.argv[1]
    output_dir = sys.argv[2]
    start_sec = int(sys.argv[3])
    end_sec = int(sys.argv[4])
    if start_sec >= end_sec:
        print("起始秒必须小于结束秒")
        sys.exit(1)
    
    segments = [(start_sec, end_sec)]
    extract_clips_with_audio(input_video, output_dir, segments)